#!/bin/bash

# =============================================================================
# 金融数据MCP服务启动脚本
# =============================================================================

set -e  # 遇到错误立即退出

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 日志函数
log_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

log_debug() {
    echo -e "${BLUE}[DEBUG]${NC} $1"
}

# 显示帮助信息
show_help() {
    cat << EOF
金融数据MCP服务启动脚本

用法: $0 [选项] [命令]

命令:
    start           启动所有服务 (默认)
    stop            停止所有服务
    restart         重启所有服务
    status          查看服务状态
    logs            查看服务日志
    build           构建Docker镜像
    init            初始化数据库
    clean           清理Docker资源
    backup          备份数据
    restore         恢复数据
    health          健康检查
    shell           进入服务容器

选项:
    -h, --help      显示此帮助信息
    -v, --verbose   详细输出
    -d, --dev       开发模式
    -p, --prod      生产模式
    --no-cache      构建时不使用缓存
    --force         强制执行操作

示例:
    $0 start                # 启动所有服务
    $0 start -d             # 开发模式启动
    $0 logs financial       # 查看金融服务日志
    $0 build --no-cache     # 重新构建镜像
    $0 backup               # 备份数据库

EOF
}

# 检查依赖
check_dependencies() {
    log_info "检查系统依赖..."
    
    # 检查Docker
    if ! command -v docker &> /dev/null; then
        log_error "Docker未安装，请先安装Docker"
        exit 1
    fi
    
    # 检查Docker Compose
    if ! command -v docker-compose &> /dev/null; then
        log_error "Docker Compose未安装，请先安装Docker Compose"
        exit 1
    fi
    
    # 检查Docker服务状态
    if ! docker info &> /dev/null; then
        log_error "Docker服务未运行，请启动Docker服务"
        exit 1
    fi
    
    log_info "依赖检查完成"
}

# 检查环境配置
check_environment() {
    log_info "检查环境配置..."
    
    if [ ! -f ".env" ]; then
        if [ -f ".env.example" ]; then
            log_warn ".env文件不存在，从.env.example复制"
            cp .env.example .env
            log_warn "请编辑.env文件配置正确的环境变量"
        else
            log_error ".env和.env.example文件都不存在"
            exit 1
        fi
    fi
    
    # 检查必要的环境变量
    source .env
    
    required_vars=(
        "POSTGRES_DB"
        "POSTGRES_USER" 
        "POSTGRES_PASSWORD"
        "REDIS_HOST"
        "REDIS_PORT"
    )
    
    for var in "${required_vars[@]}"; do
        if [ -z "${!var}" ]; then
            log_error "环境变量 $var 未设置"
            exit 1
        fi
    done
    
    log_info "环境配置检查完成"
}

# 构建Docker镜像
build_images() {
    log_info "构建Docker镜像..."
    
    local build_args=""
    if [ "$NO_CACHE" = "true" ]; then
        build_args="--no-cache"
    fi
    
    if [ "$VERBOSE" = "true" ]; then
        docker-compose build $build_args
    else
        docker-compose build $build_args > /dev/null 2>&1
    fi
    
    log_info "Docker镜像构建完成"
}

# 初始化数据库
init_database() {
    log_info "初始化数据库..."
    
    # 启动数据库服务
    docker-compose up -d postgres redis
    
    # 等待数据库启动
    log_info "等待数据库启动..."
    sleep 10
    
    # 检查数据库连接
    local max_attempts=30
    local attempt=1
    
    while [ $attempt -le $max_attempts ]; do
        if docker-compose exec -T postgres pg_isready -U $POSTGRES_USER -d $POSTGRES_DB > /dev/null 2>&1; then
            log_info "数据库连接成功"
            break
        fi
        
        log_debug "等待数据库连接... ($attempt/$max_attempts)"
        sleep 2
        ((attempt++))
    done
    
    if [ $attempt -gt $max_attempts ]; then
        log_error "数据库连接超时"
        exit 1
    fi
    
    # 执行初始化脚本
    if [ -f "init.sql" ]; then
        log_info "执行数据库初始化脚本..."
        docker-compose exec -T postgres psql -U $POSTGRES_USER -d $POSTGRES_DB -f /docker-entrypoint-initdb.d/init.sql
        log_info "数据库初始化完成"
    else
        log_warn "init.sql文件不存在，跳过数据库初始化"
    fi
}

# 启动服务
start_services() {
    log_info "启动金融数据MCP服务..."
    
    # 检查是否需要初始化数据库
    if ! docker-compose exec -T postgres psql -U $POSTGRES_USER -d $POSTGRES_DB -c "SELECT 1 FROM system_config WHERE config_key='database.initialized' LIMIT 1;" > /dev/null 2>&1; then
        log_info "检测到数据库未初始化，开始初始化..."
        init_database
    fi
    
    # 启动所有服务
    if [ "$DEV_MODE" = "true" ]; then
        log_info "开发模式启动..."
        docker-compose -f docker-compose.yml -f docker-compose.dev.yml up -d
    else
        docker-compose up -d
    fi
    
    # 等待服务启动
    log_info "等待服务启动..."
    sleep 5
    
    # 健康检查
    check_health
    
    log_info "服务启动完成"
    show_service_info
}

# 停止服务
stop_services() {
    log_info "停止金融数据MCP服务..."
    docker-compose down
    log_info "服务已停止"
}

# 重启服务
restart_services() {
    log_info "重启金融数据MCP服务..."
    stop_services
    sleep 2
    start_services
}

# 查看服务状态
show_status() {
    log_info "服务状态:"
    docker-compose ps
    
    echo ""
    log_info "容器资源使用情况:"
    docker stats --no-stream --format "table {{.Container}}\t{{.CPUPerc}}\t{{.MemUsage}}\t{{.NetIO}}\t{{.BlockIO}}"
}

# 查看日志
show_logs() {
    local service=$1
    local follow=${2:-false}
    
    if [ -z "$service" ]; then
        log_info "显示所有服务日志:"
        if [ "$follow" = "true" ]; then
            docker-compose logs -f
        else
            docker-compose logs --tail=100
        fi
    else
        log_info "显示 $service 服务日志:"
        if [ "$follow" = "true" ]; then
            docker-compose logs -f $service
        else
            docker-compose logs --tail=100 $service
        fi
    fi
}

# 健康检查
check_health() {
    log_info "执行健康检查..."
    
    local services=("postgres" "redis" "financial")
    local healthy_count=0
    
    for service in "${services[@]}"; do
        local health_status=$(docker-compose ps -q $service | xargs docker inspect --format='{{.State.Health.Status}}' 2>/dev/null || echo "unknown")
        
        case $health_status in
            "healthy")
                log_info "✓ $service: 健康"
                ((healthy_count++))
                ;;
            "unhealthy")
                log_error "✗ $service: 不健康"
                ;;
            "starting")
                log_warn "⚠ $service: 启动中"
                ;;
            *)
                log_warn "? $service: 状态未知"
                ;;
        esac
    done
    
    if [ $healthy_count -eq ${#services[@]} ]; then
        log_info "所有服务健康状态良好"
        return 0
    else
        log_warn "部分服务健康状态异常"
        return 1
    fi
}

# 显示服务信息
show_service_info() {
    echo ""
    log_info "服务访问信息:"
    echo "  API网关:     http://localhost:8000"
    echo "  API文档:     http://localhost:8000/docs"
    echo "  健康检查:    http://localhost:8000/health"
    echo "  PostgreSQL:  localhost:5432"
    echo "  Redis:       localhost:6379"
    echo ""
    log_info "管理命令:"
    echo "  查看状态:    $0 status"
    echo "  查看日志:    $0 logs [service]"
    echo "  健康检查:    $0 health"
    echo "  停止服务:    $0 stop"
}

# 备份数据
backup_data() {
    log_info "备份数据库..."
    
    local backup_dir="./backups"
    local backup_file="financial_data_$(date +%Y%m%d_%H%M%S).sql"
    
    mkdir -p $backup_dir
    
    docker-compose exec -T postgres pg_dump -U $POSTGRES_USER $POSTGRES_DB > "$backup_dir/$backup_file"
    
    if [ $? -eq 0 ]; then
        log_info "数据库备份完成: $backup_dir/$backup_file"
    else
        log_error "数据库备份失败"
        exit 1
    fi
}

# 恢复数据
restore_data() {
    local backup_file=$1
    
    if [ -z "$backup_file" ]; then
        log_error "请指定备份文件路径"
        exit 1
    fi
    
    if [ ! -f "$backup_file" ]; then
        log_error "备份文件不存在: $backup_file"
        exit 1
    fi
    
    log_warn "警告: 此操作将覆盖现有数据库数据"
    read -p "确认继续? (y/N): " -n 1 -r
    echo
    
    if [[ ! $REPLY =~ ^[Yy]$ ]]; then
        log_info "操作已取消"
        exit 0
    fi
    
    log_info "恢复数据库..."
    
    # 停止应用服务
    docker-compose stop financial
    
    # 恢复数据
    docker-compose exec -T postgres psql -U $POSTGRES_USER $POSTGRES_DB < "$backup_file"
    
    if [ $? -eq 0 ]; then
        log_info "数据库恢复完成"
        # 重启应用服务
        docker-compose start financial
    else
        log_error "数据库恢复失败"
        exit 1
    fi
}

# 清理Docker资源
clean_docker() {
    log_warn "清理Docker资源..."
    
    if [ "$FORCE" = "true" ]; then
        docker-compose down -v --remove-orphans
        docker system prune -f
        log_info "Docker资源清理完成"
    else
        read -p "确认清理所有Docker资源? (y/N): " -n 1 -r
        echo
        
        if [[ $REPLY =~ ^[Yy]$ ]]; then
            docker-compose down -v --remove-orphans
            docker system prune -f
            log_info "Docker资源清理完成"
        else
            log_info "操作已取消"
        fi
    fi
}

# 进入容器Shell
enter_shell() {
    local service=${1:-financial}
    
    log_info "进入 $service 容器..."
    
    if docker-compose ps -q $service > /dev/null; then
        docker-compose exec $service /bin/bash
    else
        log_error "服务 $service 未运行"
        exit 1
    fi
}

# 解析命令行参数
parse_args() {
    COMMAND="start"
    VERBOSE=false
    DEV_MODE=false
    PROD_MODE=false
    NO_CACHE=false
    FORCE=false
    
    while [[ $# -gt 0 ]]; do
        case $1 in
            -h|--help)
                show_help
                exit 0
                ;;
            -v|--verbose)
                VERBOSE=true
                shift
                ;;
            -d|--dev)
                DEV_MODE=true
                shift
                ;;
            -p|--prod)
                PROD_MODE=true
                shift
                ;;
            --no-cache)
                NO_CACHE=true
                shift
                ;;
            --force)
                FORCE=true
                shift
                ;;
            start|stop|restart|status|logs|build|init|clean|backup|restore|health|shell)
                COMMAND=$1
                shift
                ;;
            *)
                SERVICE_ARG=$1
                shift
                ;;
        esac
    done
    
    # 设置环境变量
    if [ "$DEV_MODE" = "true" ]; then
        export ENVIRONMENT=development
    elif [ "$PROD_MODE" = "true" ]; then
        export ENVIRONMENT=production
    fi
}

# 主函数
main() {
    echo "========================================"
    echo "    金融数据MCP服务管理脚本"
    echo "========================================"
    
    parse_args "$@"
    
    # 检查基本依赖
    check_dependencies
    
    # 根据命令执行相应操作
    case $COMMAND in
        start)
            check_environment
            start_services
            ;;
        stop)
            stop_services
            ;;
        restart)
            check_environment
            restart_services
            ;;
        status)
            show_status
            ;;
        logs)
            show_logs "$SERVICE_ARG" "$VERBOSE"
            ;;
        build)
            build_images
            ;;
        init)
            check_environment
            init_database
            ;;
        clean)
            clean_docker
            ;;
        backup)
            backup_data
            ;;
        restore)
            restore_data "$SERVICE_ARG"
            ;;
        health)
            check_health
            ;;
        shell)
            enter_shell "$SERVICE_ARG"
            ;;
        *)
            log_error "未知命令: $COMMAND"
            show_help
            exit 1
            ;;
    esac
}

# 信号处理
trap 'log_warn "收到中断信号，正在清理..."; exit 130' INT TERM

# 执行主函数
main "$@"